// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_CXX11_TENSOR_TENSOR_BROADCASTING_H
#define EIGEN_CXX11_TENSOR_TENSOR_BROADCASTING_H

namespace Eigen {

/** \class TensorBroadcasting
  * \ingroup CXX11_Tensor_Module
  *
  * \brief Tensor broadcasting class.
  *
  *
  */
namespace internal {
    template <typename Broadcast, typename XprType> struct traits<TensorBroadcastingOp<Broadcast, XprType>> : public traits<XprType>
    {
        typedef typename XprType::Scalar Scalar;
        typedef traits<XprType> XprTraits;
        typedef typename XprTraits::StorageKind StorageKind;
        typedef typename XprTraits::Index Index;
        typedef typename XprType::Nested Nested;
        typedef typename remove_reference<Nested>::type _Nested;
        static const int NumDimensions = XprTraits::NumDimensions;
        static const int Layout = XprTraits::Layout;
        typedef typename XprTraits::PointerType PointerType;
    };

    template <typename Broadcast, typename XprType> struct eval<TensorBroadcastingOp<Broadcast, XprType>, Eigen::Dense>
    {
        typedef const TensorBroadcastingOp<Broadcast, XprType> EIGEN_DEVICE_REF type;
    };

    template <typename Broadcast, typename XprType>
    struct nested<TensorBroadcastingOp<Broadcast, XprType>, 1, typename eval<TensorBroadcastingOp<Broadcast, XprType>>::type>
    {
        typedef TensorBroadcastingOp<Broadcast, XprType> type;
    };

    template <typename Dims> struct is_input_scalar
    {
        static const bool value = false;
    };
    template <> struct is_input_scalar<Sizes<>>
    {
        static const bool value = true;
    };
#ifndef EIGEN_EMULATE_CXX11_META_H
    template <typename std::ptrdiff_t... Indices> struct is_input_scalar<Sizes<Indices...>>
    {
        static const bool value = (Sizes<Indices...>::total_size == 1);
    };
#endif

}  // end namespace internal

template <typename Broadcast, typename XprType> class TensorBroadcastingOp : public TensorBase<TensorBroadcastingOp<Broadcast, XprType>, ReadOnlyAccessors>
{
public:
    typedef typename Eigen::internal::traits<TensorBroadcastingOp>::Scalar Scalar;
    typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
    typedef typename XprType::CoeffReturnType CoeffReturnType;
    typedef typename Eigen::internal::nested<TensorBroadcastingOp>::type Nested;
    typedef typename Eigen::internal::traits<TensorBroadcastingOp>::StorageKind StorageKind;
    typedef typename Eigen::internal::traits<TensorBroadcastingOp>::Index Index;

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBroadcastingOp(const XprType& expr, const Broadcast& broadcast) : m_xpr(expr), m_broadcast(broadcast) {}

    EIGEN_DEVICE_FUNC
    const Broadcast& broadcast() const { return m_broadcast; }

    EIGEN_DEVICE_FUNC
    const typename internal::remove_all<typename XprType::Nested>::type& expression() const { return m_xpr; }

protected:
    typename XprType::Nested m_xpr;
    const Broadcast m_broadcast;
};

// Eval as rvalue
template <typename Broadcast, typename ArgType, typename Device> struct TensorEvaluator<const TensorBroadcastingOp<Broadcast, ArgType>, Device>
{
    typedef TensorBroadcastingOp<Broadcast, ArgType> XprType;
    typedef typename XprType::Index Index;
    static const int NumDims = internal::array_size<typename TensorEvaluator<ArgType, Device>::Dimensions>::value;
    typedef DSizes<Index, NumDims> Dimensions;
    typedef typename XprType::Scalar Scalar;
    typedef typename TensorEvaluator<ArgType, Device>::Dimensions InputDimensions;
    typedef typename XprType::CoeffReturnType CoeffReturnType;
    typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
    static const int PacketSize = PacketType<CoeffReturnType, Device>::size;

protected:  //  all the non-static fields must have the same access control, otherwise the TensorEvaluator wont be standard layout;
    bool isCopy, nByOne, oneByN;

public:
    typedef StorageMemory<CoeffReturnType, Device> Storage;
    typedef typename Storage::Type EvaluatorPointerType;

    enum
    {
        IsAligned = TensorEvaluator<ArgType, Device>::IsAligned,
        PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
        BlockAccess = TensorEvaluator<ArgType, Device>::BlockAccess,
        PreferBlockAccess = true,
        Layout = TensorEvaluator<ArgType, Device>::Layout,
        RawAccess = false
    };

    typedef typename internal::remove_const<Scalar>::type ScalarNoConst;

    // We do block based broadcasting using a trick with 2x tensor rank and 0
    // strides. See block method implementation for details.
    typedef DSizes<Index, 2 * NumDims> BroadcastDimensions;

    //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
    typedef internal::TensorBlockDescriptor<NumDims, Index> TensorBlockDesc;
    typedef internal::TensorBlockScratchAllocator<Device> TensorBlockScratch;

    typedef typename TensorEvaluator<const ArgType, Device>::TensorBlock ArgTensorBlock;

    typedef typename internal::TensorMaterializedBlock<ScalarNoConst, NumDims, Layout, Index> TensorBlock;
    //===--------------------------------------------------------------------===//

    EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
        : isCopy(false), nByOne(false), oneByN(false), m_device(device), m_broadcast(op.broadcast()), m_impl(op.expression(), device)
    {
        // The broadcasting op doesn't change the rank of the tensor. One can't broadcast a scalar
        // and store the result in a scalar. Instead one should reshape the scalar into a a N-D
        // tensor with N >= 1 of 1 element first and then broadcast.
        EIGEN_STATIC_ASSERT((NumDims > 0), YOU_MADE_A_PROGRAMMING_MISTAKE);
        const InputDimensions& input_dims = m_impl.dimensions();
        isCopy = true;
        for (int i = 0; i < NumDims; ++i)
        {
            eigen_assert(input_dims[i] > 0);
            m_dimensions[i] = input_dims[i] * m_broadcast[i];
            if (m_broadcast[i] != 1)
            {
                isCopy = false;
            }
        }

        if (static_cast<int>(Layout) == static_cast<int>(ColMajor))
        {
            m_inputStrides[0] = 1;
            m_outputStrides[0] = 1;
            for (int i = 1; i < NumDims; ++i)
            {
                m_inputStrides[i] = m_inputStrides[i - 1] * input_dims[i - 1];
                m_outputStrides[i] = m_outputStrides[i - 1] * m_dimensions[i - 1];
            }
        }
        else
        {
            m_inputStrides[NumDims - 1] = 1;
            m_outputStrides[NumDims - 1] = 1;
            for (int i = NumDims - 2; i >= 0; --i)
            {
                m_inputStrides[i] = m_inputStrides[i + 1] * input_dims[i + 1];
                m_outputStrides[i] = m_outputStrides[i + 1] * m_dimensions[i + 1];
            }
        }

        if (input_dims[0] == 1)
        {
            oneByN = true;
            for (int i = 1; i < NumDims; ++i)
            {
                if (m_broadcast[i] != 1)
                {
                    oneByN = false;
                    break;
                }
            }
        }
        else if (input_dims[NumDims - 1] == 1)
        {
            nByOne = true;
            for (int i = 0; i < NumDims - 1; ++i)
            {
                if (m_broadcast[i] != 1)
                {
                    nByOne = false;
                    break;
                }
            }
        }

        // Handle special format like NCHW, its input shape is '[1, N..., 1]' and
        // broadcast shape is '[N, 1..., N]'
        if (!oneByN && !nByOne)
        {
            if (input_dims[0] == 1 && input_dims[NumDims - 1] == 1 && NumDims > 2)
            {
                nByOne = true;
                oneByN = true;
                for (int i = 1; i < NumDims - 1; ++i)
                {
                    if (m_broadcast[i] != 1)
                    {
                        nByOne = false;
                        oneByN = false;
                        break;
                    }
                }
            }
        }
    }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }

    EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType)
    {
        m_impl.evalSubExprsIfNeeded(NULL);
        return true;
    }

#ifdef EIGEN_USE_THREADS
    template <typename EvalSubExprsCallback> EIGEN_STRONG_INLINE void evalSubExprsIfNeededAsync(EvaluatorPointerType, EvalSubExprsCallback done)
    {
        m_impl.evalSubExprsIfNeededAsync(nullptr, [done](bool) { done(true); });
    }
#endif  // EIGEN_USE_THREADS

    EIGEN_STRONG_INLINE void cleanup() { m_impl.cleanup(); }

    EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE CoeffReturnType coeff(Index index) const
    {
        if (internal::is_input_scalar<typename internal::remove_all<InputDimensions>::type>::value)
        {
            return m_impl.coeff(0);
        }

        if (static_cast<int>(Layout) == static_cast<int>(ColMajor))
        {
            if (isCopy)
            {
                return m_impl.coeff(index);
            }
            else
            {
                return coeffColMajor(index);
            }
        }
        else
        {
            if (isCopy)
            {
                return m_impl.coeff(index);
            }
            else
            {
                return coeffRowMajor(index);
            }
        }
    }

    // TODO: attempt to speed this up. The integer divisions and modulo are slow
    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index indexColMajor(Index index) const
    {
        Index inputIndex = 0;
        EIGEN_UNROLL_LOOP
        for (int i = NumDims - 1; i > 0; --i)
        {
            const Index idx = index / m_outputStrides[i];
            if (internal::index_statically_eq<Broadcast>(i, 1))
            {
                eigen_assert(idx < m_impl.dimensions()[i]);
                inputIndex += idx * m_inputStrides[i];
            }
            else
            {
                if (internal::index_statically_eq<InputDimensions>(i, 1))
                {
                    eigen_assert(idx % m_impl.dimensions()[i] == 0);
                }
                else
                {
                    inputIndex += (idx % m_impl.dimensions()[i]) * m_inputStrides[i];
                }
            }
            index -= idx * m_outputStrides[i];
        }
        if (internal::index_statically_eq<Broadcast>(0, 1))
        {
            eigen_assert(index < m_impl.dimensions()[0]);
            inputIndex += index;
        }
        else
        {
            if (internal::index_statically_eq<InputDimensions>(0, 1))
            {
                eigen_assert(index % m_impl.dimensions()[0] == 0);
            }
            else
            {
                inputIndex += (index % m_impl.dimensions()[0]);
            }
        }
        return inputIndex;
    }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeffColMajor(Index index) const { return m_impl.coeff(indexColMajor(index)); }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index indexRowMajor(Index index) const
    {
        Index inputIndex = 0;
        EIGEN_UNROLL_LOOP
        for (int i = 0; i < NumDims - 1; ++i)
        {
            const Index idx = index / m_outputStrides[i];
            if (internal::index_statically_eq<Broadcast>(i, 1))
            {
                eigen_assert(idx < m_impl.dimensions()[i]);
                inputIndex += idx * m_inputStrides[i];
            }
            else
            {
                if (internal::index_statically_eq<InputDimensions>(i, 1))
                {
                    eigen_assert(idx % m_impl.dimensions()[i] == 0);
                }
                else
                {
                    inputIndex += (idx % m_impl.dimensions()[i]) * m_inputStrides[i];
                }
            }
            index -= idx * m_outputStrides[i];
        }
        if (internal::index_statically_eq<Broadcast>(NumDims - 1, 1))
        {
            eigen_assert(index < m_impl.dimensions()[NumDims - 1]);
            inputIndex += index;
        }
        else
        {
            if (internal::index_statically_eq<InputDimensions>(NumDims - 1, 1))
            {
                eigen_assert(index % m_impl.dimensions()[NumDims - 1] == 0);
            }
            else
            {
                inputIndex += (index % m_impl.dimensions()[NumDims - 1]);
            }
        }
        return inputIndex;
    }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeffRowMajor(Index index) const { return m_impl.coeff(indexRowMajor(index)); }

    template <int LoadMode> EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE PacketReturnType packet(Index index) const
    {
        if (internal::is_input_scalar<typename internal::remove_all<InputDimensions>::type>::value)
        {
            return internal::pset1<PacketReturnType>(m_impl.coeff(0));
        }

        if (static_cast<int>(Layout) == static_cast<int>(ColMajor))
        {
            if (isCopy)
            {
#ifdef EIGEN_GPU_COMPILE_PHASE
                // See PR 437: on NVIDIA P100 and K20m we observed a x3-4 speed up by enforcing
                // unaligned loads here. The reason is unclear though.
                return m_impl.template packet<Unaligned>(index);
#else
                return m_impl.template packet<LoadMode>(index);
#endif
            }
            else if (oneByN && !nByOne)
            {
                return packetNByOne<LoadMode>(index);
            }
            else if (!oneByN && nByOne)
            {
                return packetOneByN<LoadMode>(index);
            }
            else if (oneByN && nByOne)
            {
                return packetOneByNByOne<LoadMode>(index);
            }
            else
            {
                return packetColMajor<LoadMode>(index);
            }
        }
        else
        {
            if (isCopy)
            {
#ifdef EIGEN_GPU_COMPILE_PHASE
                // See above.
                return m_impl.template packet<Unaligned>(index);
#else
                return m_impl.template packet<LoadMode>(index);
#endif
            }
            else if (oneByN && !nByOne)
            {
                return packetOneByN<LoadMode>(index);
            }
            else if (!oneByN && nByOne)
            {
                return packetNByOne<LoadMode>(index);
            }
            else if (oneByN && nByOne)
            {
                return packetOneByNByOne<LoadMode>(index);
            }
            else
            {
                return packetRowMajor<LoadMode>(index);
            }
        }
    }

    template <int LoadMode> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packetOneByNByOne(Index index) const
    {
        EIGEN_STATIC_ASSERT((PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
        eigen_assert(index + PacketSize - 1 < dimensions().TotalSize());

        EIGEN_ALIGN_MAX typename internal::remove_const<CoeffReturnType>::type values[PacketSize];
        Index startDim, endDim;
        Index inputIndex, outputOffset, batchedIndex;

        if (static_cast<int>(Layout) == static_cast<int>(ColMajor))
        {
            startDim = NumDims - 1;
            endDim = 1;
        }
        else
        {
            startDim = 0;
            endDim = NumDims - 2;
        }

        batchedIndex = index % m_outputStrides[startDim];
        inputIndex = batchedIndex / m_outputStrides[endDim];
        outputOffset = batchedIndex % m_outputStrides[endDim];

        if (outputOffset + PacketSize <= m_outputStrides[endDim])
        {
            values[0] = m_impl.coeff(inputIndex);
            return internal::pload1<PacketReturnType>(values);
        }
        else
        {
            EIGEN_UNROLL_LOOP
            for (int i = 0, cur = 0; i < PacketSize; ++i, ++cur)
            {
                if (outputOffset + cur < m_outputStrides[endDim])
                {
                    values[i] = m_impl.coeff(inputIndex);
                }
                else
                {
                    ++inputIndex;
                    inputIndex = (inputIndex == m_inputStrides[startDim] ? 0 : inputIndex);
                    values[i] = m_impl.coeff(inputIndex);
                    outputOffset = 0;
                    cur = 0;
                }
            }
            return internal::pload<PacketReturnType>(values);
        }
    }

    template <int LoadMode> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packetOneByN(Index index) const
    {
        EIGEN_STATIC_ASSERT((PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
        eigen_assert(index + PacketSize - 1 < dimensions().TotalSize());

        Index dim, inputIndex;

        if (static_cast<int>(Layout) == static_cast<int>(ColMajor))
        {
            dim = NumDims - 1;
        }
        else
        {
            dim = 0;
        }

        inputIndex = index % m_inputStrides[dim];
        if (inputIndex + PacketSize <= m_inputStrides[dim])
        {
            return m_impl.template packet<Unaligned>(inputIndex);
        }
        else
        {
            EIGEN_ALIGN_MAX typename internal::remove_const<CoeffReturnType>::type values[PacketSize];
            EIGEN_UNROLL_LOOP
            for (int i = 0; i < PacketSize; ++i)
            {
                if (inputIndex > m_inputStrides[dim] - 1)
                {
                    inputIndex = 0;
                }
                values[i] = m_impl.coeff(inputIndex++);
            }
            return internal::pload<PacketReturnType>(values);
        }
    }

    template <int LoadMode> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packetNByOne(Index index) const
    {
        EIGEN_STATIC_ASSERT((PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
        eigen_assert(index + PacketSize - 1 < dimensions().TotalSize());

        EIGEN_ALIGN_MAX typename internal::remove_const<CoeffReturnType>::type values[PacketSize];
        Index dim, inputIndex, outputOffset;

        if (static_cast<int>(Layout) == static_cast<int>(ColMajor))
        {
            dim = 1;
        }
        else
        {
            dim = NumDims - 2;
        }

        inputIndex = index / m_outputStrides[dim];
        outputOffset = index % m_outputStrides[dim];
        if (outputOffset + PacketSize <= m_outputStrides[dim])
        {
            values[0] = m_impl.coeff(inputIndex);
            return internal::pload1<PacketReturnType>(values);
        }
        else
        {
            EIGEN_UNROLL_LOOP
            for (int i = 0, cur = 0; i < PacketSize; ++i, ++cur)
            {
                if (outputOffset + cur < m_outputStrides[dim])
                {
                    values[i] = m_impl.coeff(inputIndex);
                }
                else
                {
                    values[i] = m_impl.coeff(++inputIndex);
                    outputOffset = 0;
                    cur = 0;
                }
            }
            return internal::pload<PacketReturnType>(values);
        }
    }

    // Ignore the LoadMode and always use unaligned loads since we can't guarantee
    // the alignment at compile time.
    template <int LoadMode> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packetColMajor(Index index) const
    {
        EIGEN_STATIC_ASSERT((PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
        eigen_assert(index + PacketSize - 1 < dimensions().TotalSize());

        const Index originalIndex = index;

        Index inputIndex = 0;
        EIGEN_UNROLL_LOOP
        for (int i = NumDims - 1; i > 0; --i)
        {
            const Index idx = index / m_outputStrides[i];
            if (internal::index_statically_eq<Broadcast>(i, 1))
            {
                eigen_assert(idx < m_impl.dimensions()[i]);
                inputIndex += idx * m_inputStrides[i];
            }
            else
            {
                if (internal::index_statically_eq<InputDimensions>(i, 1))
                {
                    eigen_assert(idx % m_impl.dimensions()[i] == 0);
                }
                else
                {
                    inputIndex += (idx % m_impl.dimensions()[i]) * m_inputStrides[i];
                }
            }
            index -= idx * m_outputStrides[i];
        }
        Index innermostLoc;
        if (internal::index_statically_eq<Broadcast>(0, 1))
        {
            eigen_assert(index < m_impl.dimensions()[0]);
            innermostLoc = index;
        }
        else
        {
            if (internal::index_statically_eq<InputDimensions>(0, 1))
            {
                eigen_assert(index % m_impl.dimensions()[0] == 0);
                innermostLoc = 0;
            }
            else
            {
                innermostLoc = index % m_impl.dimensions()[0];
            }
        }
        inputIndex += innermostLoc;

        // Todo: this could be extended to the second dimension if we're not
        // broadcasting alongside the first dimension, and so on.
        if (innermostLoc + PacketSize <= m_impl.dimensions()[0])
        {
            return m_impl.template packet<Unaligned>(inputIndex);
        }
        else
        {
            EIGEN_ALIGN_MAX typename internal::remove_const<CoeffReturnType>::type values[PacketSize];
            values[0] = m_impl.coeff(inputIndex);
            EIGEN_UNROLL_LOOP
            for (int i = 1; i < PacketSize; ++i)
            {
                if (innermostLoc + i < m_impl.dimensions()[0])
                {
                    values[i] = m_impl.coeff(inputIndex + i);
                }
                else
                {
                    values[i] = coeffColMajor(originalIndex + i);
                }
            }
            PacketReturnType rslt = internal::pload<PacketReturnType>(values);
            return rslt;
        }
    }

    template <int LoadMode> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packetRowMajor(Index index) const
    {
        EIGEN_STATIC_ASSERT((PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
        eigen_assert(index + PacketSize - 1 < dimensions().TotalSize());

        const Index originalIndex = index;

        Index inputIndex = 0;
        EIGEN_UNROLL_LOOP
        for (int i = 0; i < NumDims - 1; ++i)
        {
            const Index idx = index / m_outputStrides[i];
            if (internal::index_statically_eq<Broadcast>(i, 1))
            {
                eigen_assert(idx < m_impl.dimensions()[i]);
                inputIndex += idx * m_inputStrides[i];
            }
            else
            {
                if (internal::index_statically_eq<InputDimensions>(i, 1))
                {
                    eigen_assert(idx % m_impl.dimensions()[i] == 0);
                }
                else
                {
                    inputIndex += (idx % m_impl.dimensions()[i]) * m_inputStrides[i];
                }
            }
            index -= idx * m_outputStrides[i];
        }
        Index innermostLoc;
        if (internal::index_statically_eq<Broadcast>(NumDims - 1, 1))
        {
            eigen_assert(index < m_impl.dimensions()[NumDims - 1]);
            innermostLoc = index;
        }
        else
        {
            if (internal::index_statically_eq<InputDimensions>(NumDims - 1, 1))
            {
                eigen_assert(index % m_impl.dimensions()[NumDims - 1] == 0);
                innermostLoc = 0;
            }
            else
            {
                innermostLoc = index % m_impl.dimensions()[NumDims - 1];
            }
        }
        inputIndex += innermostLoc;

        // Todo: this could be extended to the second dimension if we're not
        // broadcasting alongside the first dimension, and so on.
        if (innermostLoc + PacketSize <= m_impl.dimensions()[NumDims - 1])
        {
            return m_impl.template packet<Unaligned>(inputIndex);
        }
        else
        {
            EIGEN_ALIGN_MAX typename internal::remove_const<CoeffReturnType>::type values[PacketSize];
            values[0] = m_impl.coeff(inputIndex);
            EIGEN_UNROLL_LOOP
            for (int i = 1; i < PacketSize; ++i)
            {
                if (innermostLoc + i < m_impl.dimensions()[NumDims - 1])
                {
                    values[i] = m_impl.coeff(inputIndex + i);
                }
                else
                {
                    values[i] = coeffRowMajor(originalIndex + i);
                }
            }
            PacketReturnType rslt = internal::pload<PacketReturnType>(values);
            return rslt;
        }
    }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const
    {
        double compute_cost = TensorOpCost::AddCost<Index>();
        if (!isCopy && NumDims > 0)
        {
            EIGEN_UNROLL_LOOP
            for (int i = NumDims - 1; i > 0; --i)
            {
                compute_cost += TensorOpCost::DivCost<Index>();
                if (internal::index_statically_eq<Broadcast>(i, 1))
                {
                    compute_cost += TensorOpCost::MulCost<Index>() + TensorOpCost::AddCost<Index>();
                }
                else
                {
                    if (!internal::index_statically_eq<InputDimensions>(i, 1))
                    {
                        compute_cost += TensorOpCost::MulCost<Index>() + TensorOpCost::ModCost<Index>() + TensorOpCost::AddCost<Index>();
                    }
                }
                compute_cost += TensorOpCost::MulCost<Index>() + TensorOpCost::AddCost<Index>();
            }
        }
        return m_impl.costPerCoeff(vectorized) + TensorOpCost(0, 0, compute_cost, vectorized, PacketSize);
    }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE internal::TensorBlockResourceRequirements getResourceRequirements() const
    {
        // TODO(wuke): Targeting L1 size is 30% faster than targeting L{-1} on large
        // tensors. But this might need further tuning.
        const size_t target_size = m_device.firstLevelCacheSize();
        return internal::TensorBlockResourceRequirements::merge(m_impl.getResourceRequirements(),
                                                                internal::TensorBlockResourceRequirements::skewed<Scalar>(target_size));
    }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlock block(TensorBlockDesc& desc, TensorBlockScratch& scratch, bool /*root_of_expr_ast*/ = false) const
    {
        BlockBroadcastingParams params = blockBroadcastingParams(desc);

        if (params.inner_dim_size == 0 || params.bcast_dim_size == 0)
        {
            return emptyBlock();
        }

        // Prepare storage for the materialized broadcasting result.
        const typename TensorBlock::Storage block_storage = TensorBlock::prepareStorage(desc, scratch);
        ScalarNoConst* materialized_output = block_storage.data();

        // We potentially will need to materialize input blocks.
        size_t materialized_input_size = 0;
        ScalarNoConst* materialized_input = NULL;

        // Initialize block broadcating iterator state for outer dimensions (outer
        // with regard to bcast dimension). Dimension in this array are always in
        // inner_most -> outer_most order (col major layout).
        array<BlockBroadcastingIteratorState, NumDims> it;
        int idx = 0;

        for (int i = params.inner_dim_count + 1; i < NumDims; ++i)
        {
            const Index dim = IsColMajor ? i : NumDims - 1 - i;
            it[idx].size = params.output_dims[dim];
            it[idx].count = 0;
            it[idx].output_stride = m_outputStrides[dim];
            it[idx].output_span = it[idx].output_stride * (it[idx].size - 1);
            idx++;
        }

        // Write output into the beginning of `materialized_output`.
        Index output_offset = 0;

        // We will fill output block by broadcasting along the bcast dim, and
        // iterating over outer dimension.
        const Index output_size = NumDims == 0 ? 1 : params.output_dims.TotalSize();

        for (Index num_output_coeffs = 0; num_output_coeffs < output_size;)
        {
            ScalarNoConst* bcast_output = materialized_output + num_output_coeffs;
            Index bcast_offset = desc.offset() + output_offset;

            // Broadcast along the bcast dimension.
            num_output_coeffs += BroadcastBlockAlongBcastDim(params, bcast_offset, scratch, bcast_output, &materialized_input, &materialized_input_size);

            // Switch to the next outer dimension.
            for (int j = 0; j < idx; ++j)
            {
                if (++it[j].count < it[j].size)
                {
                    output_offset += it[j].output_stride;
                    break;
                }
                it[j].count = 0;
                output_offset -= it[j].output_span;
            }
        }

        return block_storage.AsTensorMaterializedBlock();
    }

    EIGEN_DEVICE_FUNC EvaluatorPointerType data() const { return NULL; }

    const TensorEvaluator<ArgType, Device>& impl() const { return m_impl; }

    Broadcast functor() const { return m_broadcast; }
#ifdef EIGEN_USE_SYCL
    // binding placeholder accessors to a command group handler for SYCL
    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void bind(cl::sycl::handler& cgh) const { m_impl.bind(cgh); }
#endif
private:
    static const bool IsColMajor = static_cast<int>(Layout) == static_cast<int>(ColMajor);

    // We will build a general case block broadcasting on top of broadcasting
    // primitive that will do broadcasting only for the inner dimension(s) along
    // the first dimension smaller than the input size (it's called `bcast_dim`).
    //
    // Example:
    //           dim:  0  1  2   (ColMajor)
    //    input size: [9, 3, 6]
    //    block size: [9, 2, 6]
    //
    // We will compute broadcasted block by iterating over the outer dimensions
    // before `bcast_dim` (only dimension `2` in this example) and computing
    // broadcasts along the `bcast_dim` (dimension `1` in this example).

    // BlockBroadcastingParams holds precomputed parameters for broadcasting a
    // single block along the broadcasting dimension. Sizes and strides along the
    // `bcast_dim` might be invalid, they will be adjusted later in
    // `BroadcastBlockAlongBcastDim`.
    struct BlockBroadcastingParams
    {
        Dimensions input_dims;      // input expression dimensions
        Dimensions output_dims;     // output block sizes
        Dimensions output_strides;  // output block strides

        int inner_dim_count;   // count inner dimensions matching in size
        int bcast_dim;         // broadcasting dimension index
        Index bcast_dim_size;  // broadcasting dimension size
        Index inner_dim_size;  // inner dimensions size

        // Block sizes and strides for the input block where all dimensions before
        // `bcast_dim` are equal to `1`.
        Dimensions input_block_sizes;
        Dimensions input_block_strides;

        // Block sizes and strides for blocks with extra dimensions and strides `0`.
        BroadcastDimensions bcast_block_sizes;
        BroadcastDimensions bcast_block_strides;
        BroadcastDimensions bcast_input_strides;
    };

    struct BlockBroadcastingIteratorState
    {
        Index size;
        Index count;
        Index output_stride;
        Index output_span;
    };

    EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE BlockBroadcastingParams blockBroadcastingParams(TensorBlockDesc& desc) const
    {
        BlockBroadcastingParams params;

        params.input_dims = Dimensions(m_impl.dimensions());

        // Output block sizes and strides.
        params.output_dims = desc.dimensions();
        params.output_strides = internal::strides<Layout>(params.output_dims);

        // Find the broadcasting dimension (first dimension with output size smaller
        // that the input size).
        params.bcast_dim = 0;
        params.bcast_dim_size = 1;
        params.inner_dim_size = 1;

        // Count the number of inner dimensions that have the same size in the block
        // and in the broadcast expression.
        params.inner_dim_count = 0;

        for (int i = 0; i < NumDims; ++i)
        {
            const int dim = IsColMajor ? i : NumDims - i - 1;

            if (params.output_dims[dim] == m_dimensions[dim])
            {
                params.inner_dim_size *= params.output_dims[dim];
                ++params.inner_dim_count;
                continue;
            }

            // First non-matching dimension is the broadcasting dimension.
            eigen_assert(params.output_dims[dim] < m_dimensions[dim]);
            params.bcast_dim = dim;
            params.bcast_dim_size = params.output_dims[dim];
            break;
        }

        // Calculate the input block size for looking into the input.
        for (int i = 0; i < params.inner_dim_count; ++i)
        {
            const int dim = IsColMajor ? i : NumDims - i - 1;
            params.input_block_sizes[dim] = params.input_dims[dim];
        }
        for (int i = params.inner_dim_count; i < NumDims; ++i)
        {
            const int dim = IsColMajor ? i : NumDims - i - 1;
            params.input_block_sizes[dim] = 1;
        }
        params.input_block_strides = internal::strides<Layout>(params.input_block_sizes);

        // Broadcast with the 0-stride trick: Create 1 extra dim for each
        // broadcast, set the input stride to 0.
        //
        // When ColMajor:
        //
        // - bcast_block_sizes:
        //   [d_0, b_0, d_1, b_1, ...]
        //
        // - bcast_block_strides:
        //   [output_block_strides[0], output_block_strides[0] * d_0,
        //    output_block_strides[1], output_block_strides[1] * d_1,
        //   ...]
        //
        // - bcast_input_strides:
        //   [input_block_strides[0], 0,
        //    input_block_strides[1], 0,
        //   ...].
        //
        for (int i = 0; i < params.inner_dim_count; ++i)
        {
            const int dim = IsColMajor ? i : NumDims - i - 1;

            const int copy_dim = IsColMajor ? 2 * i : 2 * NumDims - 2 * i - 1;
            const int broadcast_dim = IsColMajor ? copy_dim + 1 : copy_dim - 1;

            params.bcast_block_sizes[copy_dim] = params.input_dims[dim];
            params.bcast_block_sizes[broadcast_dim] = m_broadcast[dim];
            params.bcast_block_strides[copy_dim] = params.output_strides[dim];
            params.bcast_block_strides[broadcast_dim] = params.output_strides[dim] * params.input_dims[dim];
            params.bcast_input_strides[copy_dim] = params.input_block_strides[dim];
            params.bcast_input_strides[broadcast_dim] = 0;
        }

        for (int i = 2 * params.inner_dim_count; i < 2 * NumDims; ++i)
        {
            const int dim = IsColMajor ? i : 2 * NumDims - i - 1;
            params.bcast_block_sizes[dim] = 1;
            params.bcast_block_strides[dim] = 0;
            params.bcast_input_strides[dim] = 0;
        }

        return params;
    }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlock emptyBlock() const
    {
        DSizes<Index, NumDims> dimensions;
        for (int i = 0; i < NumDims; ++i) dimensions[i] = 0;
        return TensorBlock(internal::TensorBlockKind::kView, NULL, dimensions);
    }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index BroadcastBlockAlongBcastDim(BlockBroadcastingParams params,
                                                                            Index bcast_offset,
                                                                            TensorBlockScratch& scratch,
                                                                            ScalarNoConst* materialized_output,
                                                                            ScalarNoConst** materialized_input,
                                                                            size_t* materialized_input_size) const
    {
        if (params.bcast_dim_size == 1)
        {
            // We just need one block read using the ready-set values above.
            return BroadcastBlock(params.input_block_sizes,
                                  params.input_block_strides,
                                  params.bcast_block_sizes,
                                  params.bcast_block_strides,
                                  params.bcast_input_strides,
                                  bcast_offset,
                                  0,
                                  scratch,
                                  materialized_output,
                                  materialized_input,
                                  materialized_input_size);
        }
        else if (params.input_dims[params.bcast_dim] == 1)
        {
            // Broadcast bcast dimension (< NumDims) by bcast_dim_size.
            const int broadcast_bcast_dim = IsColMajor ? 2 * params.inner_dim_count + 1 : 2 * NumDims - 2 * params.inner_dim_count - 2;

            params.bcast_block_sizes[broadcast_bcast_dim] = params.bcast_dim_size;
            params.bcast_input_strides[broadcast_bcast_dim] = 0;
            params.bcast_block_strides[broadcast_bcast_dim] = params.output_strides[params.bcast_dim];

            return BroadcastBlock(params.input_block_sizes,
                                  params.input_block_strides,
                                  params.bcast_block_sizes,
                                  params.bcast_block_strides,
                                  params.bcast_input_strides,
                                  bcast_offset,
                                  0,
                                  scratch,
                                  materialized_output,
                                  materialized_input,
                                  materialized_input_size);
        }
        else
        {
            // Keep track of the total number of the coefficients written to the
            // output block.
            Index num_output_coeffs = 0;

            // The general case. Let's denote the output block as
            //
            //   x[..., a:a+bcast_dim_size, :, ..., :]
            //
            // where a:a+bcast_dim_size is a slice on the bcast_dim dimension
            // (< NumDims). We need to split the a:a+bcast_dim_size into possibly 3
            // sub-blocks:
            //
            // (1) a:b, where b is the smallest multiple of
            //     input_dims[bcast_dim_start] in [a, a+bcast_dim_size].
            //
            // (2) b:c, where c is the largest multiple of input_dims[bcast_dim_start]
            //     in [a, a+bcast_dim_size].
            //
            // (3) c:a+bcast_dim_size .
            //
            // Or, when b and c do not exist, we just need to process the whole block
            // together.

            // Find a.
            const Index bcast_dim_left_index = bcast_offset / m_outputStrides[params.bcast_dim];

            // Find b and c.
            const Index input_bcast_dim_size = params.input_dims[params.bcast_dim];

            // First multiple after a. This is b when <= bcast_dim_left_index +
            // bcast_dim_size.
            const Index first_multiple = divup<Index>(bcast_dim_left_index, input_bcast_dim_size) * input_bcast_dim_size;

            if (first_multiple <= bcast_dim_left_index + params.bcast_dim_size)
            {
                // b exists, so does c. Find it.
                const Index last_multiple = (bcast_dim_left_index + params.bcast_dim_size) / input_bcast_dim_size * input_bcast_dim_size;
                const int copy_bcast_dim = IsColMajor ? 2 * params.inner_dim_count : 2 * NumDims - 2 * params.inner_dim_count - 1;
                const int broadcast_bcast_dim = IsColMajor ? 2 * params.inner_dim_count + 1 : 2 * NumDims - 2 * params.inner_dim_count - 2;

                if (first_multiple > bcast_dim_left_index)
                {
                    const Index head_size = first_multiple - bcast_dim_left_index;
                    params.input_block_sizes[params.bcast_dim] = head_size;
                    params.bcast_block_sizes[copy_bcast_dim] = head_size;
                    params.bcast_input_strides[copy_bcast_dim] = params.input_block_strides[params.bcast_dim];
                    params.bcast_block_strides[copy_bcast_dim] = params.output_strides[params.bcast_dim];
                    params.bcast_block_sizes[broadcast_bcast_dim] = 1;
                    params.bcast_input_strides[broadcast_bcast_dim] = 0;
                    params.bcast_block_strides[broadcast_bcast_dim] = params.output_strides[params.bcast_dim] * params.input_dims[params.bcast_dim];

                    num_output_coeffs += BroadcastBlock(params.input_block_sizes,
                                                        params.input_block_strides,
                                                        params.bcast_block_sizes,
                                                        params.bcast_block_strides,
                                                        params.bcast_input_strides,
                                                        bcast_offset,
                                                        0,
                                                        scratch,
                                                        materialized_output,
                                                        materialized_input,
                                                        materialized_input_size);
                }
                if (first_multiple < last_multiple)
                {
                    params.input_block_sizes[params.bcast_dim] = input_bcast_dim_size;
                    params.bcast_block_sizes[copy_bcast_dim] = input_bcast_dim_size;
                    params.bcast_input_strides[copy_bcast_dim] = params.input_block_strides[params.bcast_dim];
                    params.bcast_block_strides[copy_bcast_dim] = params.output_strides[params.bcast_dim];
                    params.bcast_block_sizes[broadcast_bcast_dim] = (last_multiple - first_multiple) / input_bcast_dim_size;
                    params.bcast_input_strides[broadcast_bcast_dim] = 0;
                    params.bcast_block_strides[broadcast_bcast_dim] = params.output_strides[params.bcast_dim] * params.input_dims[params.bcast_dim];
                    const Index offset = (first_multiple - bcast_dim_left_index) * m_outputStrides[params.bcast_dim];

                    num_output_coeffs += BroadcastBlock(params.input_block_sizes,
                                                        params.input_block_strides,
                                                        params.bcast_block_sizes,
                                                        params.bcast_block_strides,
                                                        params.bcast_input_strides,
                                                        bcast_offset,
                                                        offset,
                                                        scratch,
                                                        materialized_output,
                                                        materialized_input,
                                                        materialized_input_size);
                }
                if (last_multiple < bcast_dim_left_index + params.bcast_dim_size)
                {
                    const Index tail_size = bcast_dim_left_index + params.bcast_dim_size - last_multiple;
                    params.input_block_sizes[params.bcast_dim] = tail_size;
                    params.bcast_block_sizes[copy_bcast_dim] = tail_size;
                    params.bcast_input_strides[copy_bcast_dim] = params.input_block_strides[params.bcast_dim];
                    params.bcast_block_strides[copy_bcast_dim] = params.output_strides[params.bcast_dim];
                    params.bcast_block_sizes[broadcast_bcast_dim] = 1;
                    params.bcast_input_strides[broadcast_bcast_dim] = 0;
                    params.bcast_block_strides[broadcast_bcast_dim] = params.output_strides[params.bcast_dim] * params.input_dims[params.bcast_dim];
                    const Index offset = (last_multiple - bcast_dim_left_index) * m_outputStrides[params.bcast_dim];

                    num_output_coeffs += BroadcastBlock(params.input_block_sizes,
                                                        params.input_block_strides,
                                                        params.bcast_block_sizes,
                                                        params.bcast_block_strides,
                                                        params.bcast_input_strides,
                                                        bcast_offset,
                                                        offset,
                                                        scratch,
                                                        materialized_output,
                                                        materialized_input,
                                                        materialized_input_size);
                }
            }
            else
            {
                // b and c do not exist.
                const int copy_bcast_dim = IsColMajor ? 2 * params.inner_dim_count : 2 * NumDims - 2 * params.inner_dim_count - 1;
                params.input_block_sizes[params.bcast_dim] = params.bcast_dim_size;
                params.bcast_block_sizes[copy_bcast_dim] = params.bcast_dim_size;
                params.bcast_input_strides[copy_bcast_dim] = params.input_block_strides[params.bcast_dim];
                params.bcast_block_strides[copy_bcast_dim] = params.output_strides[params.bcast_dim];

                num_output_coeffs += BroadcastBlock(params.input_block_sizes,
                                                    params.input_block_strides,
                                                    params.bcast_block_sizes,
                                                    params.bcast_block_strides,
                                                    params.bcast_input_strides,
                                                    bcast_offset,
                                                    0,
                                                    scratch,
                                                    materialized_output,
                                                    materialized_input,
                                                    materialized_input_size);
            }

            return num_output_coeffs;
        }
    }

    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index BroadcastBlock(const Dimensions& input_block_sizes,
                                                               const Dimensions& input_block_strides,
                                                               const BroadcastDimensions& bcast_block_sizes,
                                                               const BroadcastDimensions& bcast_block_strides,
                                                               const BroadcastDimensions& bcast_input_strides,
                                                               Index bcast_offset,
                                                               Index offset,
                                                               TensorBlockScratch& scratch,
                                                               ScalarNoConst* materialized_output,
                                                               ScalarNoConst** materialized_input,
                                                               size_t* materialized_input_size) const
    {
        // ---------------------------------------------------------------------- //
        // Tensor block descriptor for reading block from the input.
        const Index input_offset = bcast_offset + offset;
        TensorBlockDesc input_desc(IsColMajor ? indexColMajor(input_offset) : indexRowMajor(input_offset), input_block_sizes);

        ArgTensorBlock input_block = m_impl.block(input_desc, scratch);

        // ---------------------------------------------------------------------- //
        // Materialize input block into a temporary memory buffer only if it's not
        // already available in the arg block.
        const ScalarNoConst* input_buffer = NULL;

        if (input_block.data() != NULL)
        {
            // Input block already has raw data, there is no need to materialize it.
            input_buffer = input_block.data();
        }
        else
        {
            // Otherwise we have to do block assignment into a temporary buffer.

            // Maybe reuse previously allocated buffer, or allocate a new one with a
            // scratch allocator.
            const size_t input_total_size = input_block_sizes.TotalSize();
            if (*materialized_input == NULL || *materialized_input_size < input_total_size)
            {
                *materialized_input_size = input_total_size;
                void* mem = scratch.allocate(*materialized_input_size * sizeof(Scalar));
                *materialized_input = static_cast<ScalarNoConst*>(mem);
            }

            typedef internal::TensorBlockAssignment<ScalarNoConst, NumDims, typename ArgTensorBlock::XprType, Index> TensorBlockAssignment;

            TensorBlockAssignment::Run(TensorBlockAssignment::target(input_block_sizes, input_block_strides, *materialized_input), input_block.expr());

            input_buffer = *materialized_input;
        }

        // ---------------------------------------------------------------------- //
        // Copy data from materialized input block to the materialized output, using
        // given broadcast strides (strides with zeroes).
        typedef internal::TensorBlockIO<ScalarNoConst, Index, 2 * NumDims, Layout> TensorBlockIO;

        typename TensorBlockIO::Src src(bcast_input_strides, input_buffer);
        typename TensorBlockIO::Dst dst(bcast_block_sizes, bcast_block_strides, materialized_output + offset);

        return TensorBlockIO::Copy(dst, src);
    }

protected:
    const Device EIGEN_DEVICE_REF m_device;
    const typename internal::remove_reference<Broadcast>::type m_broadcast;
    Dimensions m_dimensions;
    array<Index, NumDims> m_outputStrides;
    array<Index, NumDims> m_inputStrides;
    TensorEvaluator<ArgType, Device> m_impl;
};

}  // end namespace Eigen

#endif  // EIGEN_CXX11_TENSOR_TENSOR_BROADCASTING_H
